#include <UnitGroupManager.h>
#include "EnhancedUI.h"
#include "CFormation.h"

using namespace BWAPI;

EnhancedUI::EnhancedUI()
{
	analyzed=false;
	show_stats=false;
	show_visibility_data=false;
	show_attack_lines=true;
	show_move_paths=true;
	show_formations=true;
	show_idle=true;
	show_enemy_sight_range = false;
	show_enemy_attack_range = false;
	show_manager_assignments = false;
}

void EnhancedUI::setArbitrator(Arbitrator::Arbitrator<BWAPI::Unit*,double>* arbitrator)
{
  this->arbitrator = arbitrator;
}

void EnhancedUI::setAnalyzed(bool analyzed)
{
	this->analyzed = analyzed;
}

void EnhancedUI::update() const
{
	if (analyzed)
	{
		drawTerrain();
		drawBases();
	}
	drawProgress();
	drawUnitSymbols();
	drawEnemySymbols();

	if (show_stats)
		void drawStats();

	if (show_formations)
		drawFormations();

	if (show_visibility_data)
		drawVisibility();

	if (show_manager_assignments)
		showManagerData();
}

void EnhancedUI::drawBases() const
{
	//we will iterate through all the base locations, and draw their outlines.
	for(std::set<BWTA::BaseLocation*>::const_iterator i=BWTA::getBaseLocations().begin();i!=BWTA::getBaseLocations().end();i++)
	{
		TilePosition p=(*i)->getTilePosition();
		Position c=(*i)->getPosition();

		//draw outline of center location
		Broodwar->drawBox(CoordinateType::Map,p.x()*32,p.y()*32,p.x()*32+4*32,p.y()*32+3*32,Colors::Blue,false);

		//draw a circle at each mineral patch
		for(std::set<BWAPI::Unit*>::const_iterator j=(*i)->getMinerals().begin();j!=(*i)->getMinerals().end();j++)
		{
			Position q=(*j)->getPosition();
			Broodwar->drawCircle(CoordinateType::Map,q.x(),q.y(),30,Colors::Cyan,false);
		}

		//draw the outlines of vespene geysers
		for(std::set<BWAPI::Unit*>::const_iterator j=(*i)->getGeysers().begin();j!=(*i)->getGeysers().end();j++)
		{
			TilePosition q=(*j)->getTilePosition();
			Broodwar->drawBox(CoordinateType::Map,q.x()*32,q.y()*32,q.x()*32+4*32,q.y()*32+2*32,Colors::Orange,false);
		}

		//if this is an island expansion, draw a yellow circle around the base location
		if ((*i)->isIsland())
		{
			Broodwar->drawCircle(CoordinateType::Map,c.x(),c.y(),80,Colors::Yellow,false);
		}
	}
}

void EnhancedUI::drawTerrain() const
{
	//we will iterate through all the regions and draw the polygon outline of it in green.
	for(std::set<BWTA::Region*>::const_iterator r=BWTA::getRegions().begin();r!=BWTA::getRegions().end();r++)
	{
		BWTA::Polygon p=(*r)->getPolygon();
		for(int j=0;j<(int)p.size();j++)
		{
			Position point1=p[j];
			Position point2=p[(j+1) % p.size()];
			Broodwar->drawLine(CoordinateType::Map,point1.x(),point1.y(),point2.x(),point2.y(),Colors::Green);
		}
	}

	//we will visualize the chokepoints with red lines
	for(std::set<BWTA::Region*>::const_iterator r=BWTA::getRegions().begin();r!=BWTA::getRegions().end();r++)
	{
		for(std::set<BWTA::Chokepoint*>::const_iterator c=(*r)->getChokepoints().begin();c!=(*r)->getChokepoints().end();c++)
		{
			Position center = (*c)->getCenter();
			Broodwar->drawCircle(CoordinateType::Map, center.x(), center.y(), (*c)->getWidth()/2, Colors::Red, false);
		}
	}
}

void EnhancedUI::drawProgress() const
{
	UnitGroup constructing = SelectAll()(isBuilding).not(isCompleted);
	for each (Unit* c in constructing)
	{
		double progress = 1.0 - ((static_cast<double>(c->getRemainingBuildTime())) / c->getType().buildTime());
		drawProgressBar(c->getPosition(), progress, BWAPI::Colors::Red);
	}
	UnitGroup producing = SelectAll()(isTraining);
	for each (Unit* c in producing)
	{
		if (c->getRemainingTrainTime() > .0)
		{
			double progress = 1.0 - ((static_cast<double>(c->getRemainingTrainTime())) / c->getTrainingQueue().front().buildTime());
			drawProgressBar(c->getPosition(), progress, BWAPI::Colors::Green);
		}
	}
	UnitGroup researching = SelectAll()(isResearching);
	for each (Unit* c in researching)
	{
		if (c->getRemainingResearchTime() > .0)
		{
			double progress = 1.0 - ((static_cast<double>(c->getRemainingResearchTime())) / c->getTrainingQueue().front().buildTime());
			drawProgressBar(c->getPosition(), progress, BWAPI::Colors::Cyan);
		}
	}

	UnitGroup upgrading = SelectAll()(isUpgrading);
	for each (Unit* c in upgrading)
	{
		if (c->getRemainingUpgradeTime() > .0)
		{
			double progress = 1.0 - ((static_cast<double>(c->getRemainingUpgradeTime())) / 
				(c->getUpgrade().upgradeTimeBase()+c->getUpgrade().upgradeTimeFactor()*Broodwar->self()->getUpgradeLevel(c->getUpgrade())));
			drawProgressBar(c->getPosition(), progress, BWAPI::Colors::Yellow);
		}
	}
}

void EnhancedUI::drawProgressBar(BWAPI::Position pos, double progressFraction, BWAPI::Color innerBar) const
{
	const int width = 20, height = 4;
	const BWAPI::Color outline = BWAPI::Colors::Blue;
	const BWAPI::Color barBG = BWAPI::Color(0, 0, 170);
	int xLeft = pos.x() - width / 2, xRight = pos.x() + width / 2;
	int yTop = pos.y() - height / 2, yBottom = pos.y() + height / 2;

	//Draw outline
	Broodwar->drawLineMap(xLeft + 1, yTop, xRight - 1, yTop, outline);        //top
	Broodwar->drawLineMap(xLeft + 1, yBottom, xRight - 1, yBottom, outline);  //bottom
	Broodwar->drawLineMap(xLeft, yTop + 1, xLeft, yBottom, outline);          //left
	Broodwar->drawLineMap(xRight - 1, yTop + 1, xRight - 1, yBottom, outline);//right
	//Draw bar
	Broodwar->drawBoxMap(xLeft + 1, yTop + 1, xRight - 1, yBottom, barBG, true);
	//Draw progress bar
	const int innerWidth = (xRight - 1) - (xLeft + 1);
	int progressWidth = static_cast<int>(progressFraction * innerWidth);
	Broodwar->drawBoxMap(xLeft + 1, yTop + 1, (xLeft + 1) + progressWidth, yBottom, innerBar, true);
}

void EnhancedUI::drawUnitSymbols() const
{
	for(std::set<Unit*>::const_iterator i = Broodwar->self()->getUnits().begin(); i != Broodwar->self()->getUnits().end(); i++)
	{
		Position pos = (*i)->getPosition();
		if (show_idle && (*i)->isIdle())
		{
			Broodwar->drawCircle(CoordinateType::Map, (*i)->getPosition().x(), (*i)->getPosition().y(), 10, Colors::Yellow, false);
			//if (!(*i)->getType().isBuilding() && (Broodwar->getFrameCount()%100 == 0))
			//      Broodwar->printf("A %s is sitting idle at (%d,%d)", (*i)->getType().getName().c_str(), (*i)->getPosition().x(), (*i)->getPosition().y());
		}
		if (show_move_paths)
		{
			Position point2 = (*i)->getTargetPosition();
			Broodwar->drawLine(CoordinateType::Map, pos.x(), pos.y(), point2.x(), point2.y(), Colors::Green);
		}
		if (show_attack_lines)
		{
			if (Unit* target = (*i)->getTarget())
			{
				Position point2 = target->getPosition();
				Broodwar->drawLine(CoordinateType::Map, pos.x(), pos.y(), point2.x(), point2.y(), Colors::Red);
			}
		}
	}
}

void EnhancedUI::drawEnemySymbols() const
{
	UnitGroup enemy = SelectAllEnemy();
	for each (Unit* c in enemy)
	{
		Position pos = c->getPosition();
		if (show_enemy_sight_range)
		{
			int radius = c->getType().sightRange();
			Broodwar->drawCircle(CoordinateType::Map, pos.x(), pos.y(), radius, Colors::Grey);
		}
		if (show_enemy_attack_range)
		{
			int radius;
			radius = c->getType().airWeapon()->maxRange();
			Broodwar->drawCircle(CoordinateType::Map, pos.x(), pos.y(), radius, Colors::Cyan);
			radius = c->getType().groundWeapon()->maxRange();
			Broodwar->drawCircle(CoordinateType::Map, pos.x(), pos.y(), radius, Colors::Red);
		}
	}
}

void EnhancedUI::showManagerData() const
{
	if (arbitrator)
	{
		std::set<Unit*> units=Broodwar->self()->getUnits();
		for(std::set<Unit*>::iterator i=units.begin();i!=units.end();i++)
		{
			if (this->arbitrator->hasBid(*i))
			{
				int x=(*i)->getPosition().x();
				int y=(*i)->getPosition().y();
				std::list< std::pair< Arbitrator::Controller<BWAPI::Unit*,double>*, double> > bids=this->arbitrator->getAllBidders(*i);
				int y_off=0;
				bool first = false;
				const char activeColor = '\x07', inactiveColor = '\x16';
				char color = activeColor;
				for(std::list< std::pair< Arbitrator::Controller<BWAPI::Unit*,double>*, double> >::iterator j=bids.begin();j!=bids.end();j++)
				{
					Broodwar->drawText(CoordinateType::Map,x,y+y_off,"%c%s: %d",color,j->first->getShortName().c_str(),(int)j->second);
					y_off+=15;
					color = inactiveColor;
				}
			}
		}
	}
}

void EnhancedUI::drawStats() const
{
	std::set<Unit*> myUnits = Broodwar->self()->getUnits();
	Broodwar->drawTextScreen(5,0,"I have %d units:",myUnits.size());
	std::map<UnitType, int> unitTypeCounts;
	for(std::set<Unit*>::iterator i=myUnits.begin();i!=myUnits.end();i++)
	{
		if (unitTypeCounts.find((*i)->getType())==unitTypeCounts.end())
		{
			unitTypeCounts.insert(std::make_pair((*i)->getType(),0));
		}
		unitTypeCounts.find((*i)->getType())->second++;
	}
	int line=1;
	for(std::map<UnitType,int>::iterator i=unitTypeCounts.begin();i!=unitTypeCounts.end();i++)
	{
		Broodwar->drawTextScreen(5,16*line,"- %d %ss",(*i).second, (*i).first.getName().c_str());
		line++;
	}
}

void EnhancedUI::showPlayers()
{
	std::set<Player*> players=Broodwar->getPlayers();
	for(std::set<Player*>::iterator i=players.begin();i!=players.end();i++)
	{
		Broodwar->printf("Player [%d]: %s is in force: %s",(*i)->getID(),(*i)->getName().c_str(), (*i)->getForce()->getName().c_str());
	}
}
void EnhancedUI::showForces()
{
	std::set<Force*> forces=Broodwar->getForces();
	for(std::set<Force*>::iterator i=forces.begin();i!=forces.end();i++)
	{
		std::set<Player*> players=(*i)->getPlayers();
		Broodwar->printf("Force %s has the following players:",(*i)->getName().c_str());
		for(std::set<Player*>::iterator j=players.begin();j!=players.end();j++)
		{
			Broodwar->printf("  - Player [%d]: %s",(*j)->getID(),(*j)->getName().c_str());
		}
	}
}

void EnhancedUI::drawVisibility() const
{
	for(int x=0;x<Broodwar->mapWidth();x++)
	{
		for(int y=0;y<Broodwar->mapHeight();y++)
		{
			if (Broodwar->isExplored(x,y))
			{
				if (Broodwar->isVisible(x,y))
					Broodwar->drawDotMap(x*32+16,y*32+16,Colors::Green);
				else
					Broodwar->drawDotMap(x*32+16,y*32+16,Colors::Blue);
			}
			else
				Broodwar->drawDotMap(x*32+16,y*32+16,Colors::Red);
		}
	}
}

void EnhancedUI::drawFormations() const
{
	std::set<CFormation*> formations = CFormation::getFormations();
	for (std::set<CFormation*>::iterator i = formations.begin(); i != formations.end(); ++i)
	{
		Broodwar->drawCircle(CoordinateType::Map, (*i)->getCenter().x(), (*i)->getCenter().y(), 3, Colors::Yellow, true);
		for (int j = 0; j < (*i)->getUnits().size(); ++j)
		{
			BWAPI::Position pos = (*i)->getPosition(j);
			Broodwar->drawCircle(CoordinateType::Map, pos.x(), pos.y(), 2, Colors::Green, true);
		}
		if (Unit* target = (*i)->getTarget())
			Broodwar->drawLine(CoordinateType::Map,
			(*i)->getCenter().x(), (*i)->getCenter().y(),
			target->getPosition().x(), target->getPosition().y(),
			Colors::Red);
	}
}

bool EnhancedUI::onSendText(std::string text)
{
	if (text=="/debug")
	{
		show_manager_assignments = !show_manager_assignments;
		return false;
	} else if (text=="/show visibility")
	{
		show_visibility_data=true;
		return false;
	} else if (text=="/show players")
	{
		showPlayers();
		return false;
	} else if (text=="/show forces")
	{
		showForces();
		return false;
	} else if (text=="/show idle")
	{
		show_idle=true;
		return false;
	} else if (text=="/hide idle")
	{
		show_idle=false;
		return false;
	} else if (text=="/show enemy sight range")
	{
		show_enemy_sight_range=true;
		return false;
	} else if (text=="/hide enemy sight range")
	{
		show_enemy_sight_range=false;
		return false;
	} else if (text=="/show enemy attack range")
	{
		show_enemy_attack_range=true;
		return false;
	} else if (text=="/hide enemy attack range")
	{
		show_enemy_attack_range=false;
		return false;
	} else if (text=="/show targets")
	{
		show_attack_lines=true;
		return false;
	} else if (text=="/hide targets")
	{
		show_attack_lines=false;
		return false;
	} else if (text=="/show paths")
	{
		show_move_paths=true;
		return false;
	} else if (text=="/hide paths")
	{
		show_move_paths=false;
		return false;
	} else if (text=="/show formations")
	{
		show_formations=true;
		return false;
	} else if (text=="/hide formations")
	{
		show_formations=false;
		return false;
	}
	return true;
}